package cn.xiaopengstack.domain.agent.service.armory;

import cn.bugstack.wrench.design.framework.tree.AbstractMultiThreadStrategyRouter;
import cn.xiaopengstack.domain.agent.adapter.repository.IAgentRepository;
import cn.xiaopengstack.domain.agent.model.entity.ArmoryCommandEntity;
import cn.xiaopengstack.domain.agent.model.valobj.AiClientToolMcpVO;
import cn.xiaopengstack.domain.agent.model.valobj.McpTransportTypeEnumVO;
import cn.xiaopengstack.domain.agent.service.armory.factory.DefaultArmoryStrategyFactory;
import io.modelcontextprotocol.client.McpClient;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.client.transport.HttpClientSseClientTransport;
import io.modelcontextprotocol.client.transport.ServerParameters;
import io.modelcontextprotocol.client.transport.StdioClientTransport;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.config.BeanDefinition;
import org.springframework.beans.factory.support.BeanDefinitionBuilder;
import org.springframework.beans.factory.support.DefaultListableBeanFactory;
import org.springframework.context.ApplicationContext;

import java.time.Duration;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeoutException;

/**
 * @author jiangyangang
 * 装配支撑层
 */
@Slf4j
public abstract class AbstractArmorySupport extends AbstractMultiThreadStrategyRouter<ArmoryCommandEntity, DefaultArmoryStrategyFactory.DynamicContext, String> {

    @Resource
    protected ApplicationContext applicationContext;

    @Resource
    protected ThreadPoolExecutor threadPoolExecutor;

    @Resource
    protected IAgentRepository repository;

    @Override
    protected void multiThread(ArmoryCommandEntity requestParameter, DefaultArmoryStrategyFactory.DynamicContext dynamicContext) throws ExecutionException, InterruptedException, TimeoutException {

    }

    protected abstract String beanName(String id);

    protected abstract String dataName();

    protected synchronized <T> void registerBean(String beanName, Class<T> beanClass, T beanInstance){
        DefaultListableBeanFactory beanFactory = (DefaultListableBeanFactory) applicationContext.getAutowireCapableBeanFactory();

        // 注册Bean
        BeanDefinitionBuilder beanDefinitionBuilder = BeanDefinitionBuilder.genericBeanDefinition(beanClass, () -> beanInstance);
        BeanDefinition beanDefinition = beanDefinitionBuilder.getRawBeanDefinition();
        // 单例模式注入
        beanDefinition.setScope(BeanDefinition.SCOPE_SINGLETON);

        if (beanFactory.containsBeanDefinition(beanName)){
            beanFactory.removeBeanDefinition(beanName);
        }

        // 注册当前新的bean
        beanFactory.registerBeanDefinition(beanName, beanDefinition);

        log.info("装配支撑层，成功注册Bean:{}", beanName);
    }

    /**
     * 构造MCP客户端
     * @param mcpVO
     * @return
     */
    protected McpSyncClient getMcpClient(AiClientToolMcpVO mcpVO) {
        McpSyncClient mcpSyncClient = null;
        try {
            // 构建stdio客户端
            if (McpTransportTypeEnumVO.stdio.getType().equals(mcpVO.getTransportType())) {

                AiClientToolMcpVO.TransportConfigStdio.Stdio stdio = mcpVO.getTransportConfigStdio().getStdio().get(mcpVO.getMcpName());
                ServerParameters serverParameters = ServerParameters.builder(stdio.getCommand())
                        .args(stdio.getArgs())
                        .env(stdio.getEnv())
                        .build();
                mcpSyncClient = McpClient.sync(new StdioClientTransport(serverParameters))
                        .requestTimeout(Duration.ofMinutes(mcpVO.getRequestTimeout()))
                        .build();
                mcpSyncClient.initialize();
            }

            // 构建sse客户端
            if (McpTransportTypeEnumVO.sse.getType().equals(mcpVO.getTransportType())) {
                AiClientToolMcpVO.TransportConfigSse transportConfigSse = mcpVO.getTransportConfigSse();
                String originalBaseUri = transportConfigSse.getBaseUri();
                String baseUri;
                String sseEndpoint;
                int queryParamStartIndex = originalBaseUri.indexOf("sse");
                if (queryParamStartIndex != -1) {
                    baseUri = originalBaseUri.substring(0, queryParamStartIndex - 1);
                    sseEndpoint = originalBaseUri.substring(queryParamStartIndex - 1);
                } else {
                    baseUri = originalBaseUri;
                    sseEndpoint = transportConfigSse.getSseEndpoint();
                }

                sseEndpoint = StringUtils.isBlank(sseEndpoint) ? "/sse" : sseEndpoint;


                HttpClientSseClientTransport sseClientTransport = HttpClientSseClientTransport.builder(baseUri).sseEndpoint(sseEndpoint).build();

                mcpSyncClient = McpClient.sync(sseClientTransport).requestTimeout(Duration.ofSeconds(mcpVO.getRequestTimeout())).build();
                mcpSyncClient.initialize();
            }
        } catch (Exception e) {
            log.error("初始化MCP客户端失败!", e);
        }

        return mcpSyncClient;
    }

    /**
     * 获取Bean对象
     * @param beanName
     * @return
     */
    protected <T> T getBean(String beanName){
        return (T) applicationContext.getBean(beanName);
    }
}
